import sm9

if __name__ == '__main__':
    # 输入用户id 不可省
    idA = 'c'
    idB = 'b'

    print ("-----------------test sign and verify---------------")
    
    master_public, master_secret = sm9.setup ('sign')

    # 提取私钥用于签名
    Da = sm9.private_key_extract ('sign', master_public, master_secret, idA)
    
    # 输入待加密原始信息
    message = 'abc'

    # 进行签名
    signature = sm9.sign (master_public, Da, message)

    # 验证签名结果
    assert (sm9.verify (master_public, idA, message, signature))

    print ("\t\t\t success")

    print ("-----------------test key agreement---------------")
    
    master_public, master_secret = sm9.setup ('keyagreement')

    Da = sm9.private_key_extract ('keyagreement', master_public, master_secret, idA)
    Db = sm9.private_key_extract ('keyagreement', master_public, master_secret, idB)

    xa, Ra = sm9.generate_ephemeral (master_public, idB)
    xb, Rb = sm9.generate_ephemeral (master_public, idA)

    ska = sm9.generate_session_key (idA, idB, Ra, Rb, Da, xa, master_public, 'A', 128)
    skb = sm9.generate_session_key (idA, idB, Ra, Rb, Db, xb, master_public, 'B', 128)

    assert (ska == skb)

    print ("\t\t\t success")
    
    print ("-----------------test encrypt and decrypt---------------")

    master_public, master_secret = sm9.setup ('encrypt')

    Da = sm9.private_key_extract ('encrypt', master_public, master_secret, idA)

    message = 'abc'
    print('message:')
    print(message)

    ct = sm9.kem_dem_enc (master_public, idA, message, 32)
    print('encrypted text:')
    print(ct)

    pt = sm9.kem_dem_dec (master_public, idA, Da, ct, 32)
    print('plain text:')
    print(pt)

    assert (message == pt)

    print ("\t\t\t success")
